#![allow(dead_code, unused)]

use std::{collections::HashMap, ffi::*, ptr::{copy, copy_nonoverlapping, null, null_mut}, sync::{Arc, Mutex}};
use ash::vk::{self, Handle};
use getset::{CloneGetters, CopyGetters, Getters, MutGetters, Setters, WithSetters};
use crate::engine::njm_device::NjmDevice;

// NjmDescriptorSetLayoutBuilder
#[derive(Getters, Setters, WithSetters, MutGetters, CopyGetters, CloneGetters)]
pub struct NjmDescriptorSetLayoutBuilder<'a> {
    dev: Arc<Mutex<NjmDevice>>,
    
    #[getset(get = "pub")]
    bindings: HashMap<u32, vk::DescriptorSetLayoutBinding<'a>>,
}

impl<'a> NjmDescriptorSetLayoutBuilder<'a> {

    const DEAULT_DESCRIPTOR_COUNT: u32 = 1;

    pub fn new(dev: Arc<Mutex<NjmDevice>>) -> Self {
        return Self {
            dev: dev,
            bindings: HashMap::new(),
        };
    }

    pub fn add_binding(&mut self, 
        binding: u32,
        descriptor_type: vk::DescriptorType,
        stage_flags: vk::ShaderStageFlags,
        count: u32,
    ) -> &mut Self {
        assert!(!self.bindings.contains_key(&binding), "Binding already in use");
        let layout_binding = vk::DescriptorSetLayoutBinding::default()
            .binding(binding)
            .descriptor_type(descriptor_type)
            .descriptor_count(count)
            .stage_flags(stage_flags);
        self.bindings.insert(binding, layout_binding);
        return self;
    }

    pub fn build(&'a self) -> Box<NjmDescriptorSetLayout<'a>> {
        return Box::new(NjmDescriptorSetLayout::new(self.dev.clone(), self.bindings.clone()));
    }
}

// NjmDescriptorSetLayout
#[derive(Getters, Setters, WithSetters, MutGetters, CopyGetters, CloneGetters)]
pub struct NjmDescriptorSetLayout<'a> {
    dev: Arc<Mutex<NjmDevice>>,
    
    #[getset(get = "pub")]
    bindings: HashMap<u32, vk::DescriptorSetLayoutBinding<'a>>,

    #[getset(get = "pub")]
    descriptor_set_layout: vk::DescriptorSetLayout,
}

impl<'a> NjmDescriptorSetLayout<'a> {
    pub fn new(dev: Arc<Mutex<NjmDevice>>, 
        bindings: HashMap<u32, vk::DescriptorSetLayoutBinding<'a>>,
    ) -> Self {
        let mut set_layout_bindings: Vec<vk::DescriptorSetLayoutBinding<'_>> = vec![];
        for bind_kv in &bindings {
            set_layout_bindings.push(*bind_kv.1);
        }
        let descriptor_set_layout_info = vk::DescriptorSetLayoutCreateInfo::default()
            .bindings(&set_layout_bindings);
        let njm_dev = dev.lock().unwrap();
        let layout = unsafe { njm_dev.device().create_descriptor_set_layout(&descriptor_set_layout_info, None).unwrap() };

        return Self {
            dev: dev.clone(),
            bindings: bindings,
            descriptor_set_layout: layout,
        };
    }
}

impl Drop for NjmDescriptorSetLayout<'_> {
    fn drop(&mut self) {
        unsafe { self.dev.lock().unwrap().device().destroy_descriptor_set_layout(self.descriptor_set_layout, None) };
    }
}

// NjmDescriptorPoolBuilder
#[derive(Getters, Setters, WithSetters, MutGetters, CopyGetters, CloneGetters)]
pub struct NjmDescriptorPoolBuilder {
    dev: Arc<Mutex<NjmDevice>>,

    #[getset(get = "pub")]
    pool_sizes: Vec<vk::DescriptorPoolSize>,

    #[getset(get = "pub")]
    max_sets: u32,

    #[getset(get = "pub")]
    pool_flags: vk::DescriptorPoolCreateFlags,
}

impl NjmDescriptorPoolBuilder {
    const DEFAULT_MAX_SETS: u32 = 1000;
    
    pub fn new(dev: Arc<Mutex<NjmDevice>>) -> Self {
        return Self {
            dev: dev,
            pool_sizes: vec![],
            max_sets: Self::DEFAULT_MAX_SETS,
            pool_flags: vk::DescriptorPoolCreateFlags::default(),
        };
    }

    pub fn add_pool_size(&mut self, desc_type: vk::DescriptorType, count: u32) -> &mut Self {
        self.pool_sizes.push(
            vk::DescriptorPoolSize {
                ty: desc_type, 
                descriptor_count: count,
            }
        );
        return self;
    }

    pub fn set_pool_flags(&mut self, flags: vk::DescriptorPoolCreateFlags) -> &mut Self {
        self.pool_flags = flags;
        return self;
    }

    pub fn set_max_sets(&mut self, count: u32) -> &mut Self {
        self.max_sets = count;
        return self;
    }

    pub fn build(&self) -> Box<NjmDescriptorPool> {
        return Box::new(NjmDescriptorPool::new(
            self.dev.clone(), 
            &self.pool_sizes, 
            self.max_sets, 
            self.pool_flags
        ));
    }
}

// NjmDescriptorPool
#[derive(Getters, Setters, WithSetters, MutGetters, CopyGetters, CloneGetters)]
pub struct NjmDescriptorPool {
    dev: Arc<Mutex<NjmDevice>>,

    #[getset(get = "pub")]
    desc_pool: vk::DescriptorPool,
}

impl NjmDescriptorPool {
    
    pub fn new(
        dev: Arc<Mutex<NjmDevice>>,
        pool_sizes: &Vec<vk::DescriptorPoolSize>,
        max_sets: u32,
        pool_flags: vk::DescriptorPoolCreateFlags,
    ) -> Self {
        let desc_pool_info = vk::DescriptorPoolCreateInfo::default()
            .max_sets(max_sets)
            .pool_sizes(pool_sizes)
            .flags(pool_flags);
        let njm_dev = dev.lock().unwrap();
        let pool = unsafe { njm_dev.device().create_descriptor_pool(&desc_pool_info, None).unwrap() };
        return Self {
            dev: dev.clone(),
            desc_pool: pool,
        };
    }

    pub fn allocate_descriptor(&self, 
        layout: vk::DescriptorSetLayout,
    ) -> vk::DescriptorSet {
        let layouts = [layout];
        let alloc_info = vk::DescriptorSetAllocateInfo::default()
            .descriptor_pool(self.desc_pool)
            .set_layouts(&layouts);
        let njm_dev = self.dev.lock().unwrap();
        let desc_sets = unsafe { njm_dev.device().allocate_descriptor_sets(&alloc_info).unwrap() };
        return desc_sets[0];
    }

    pub fn free_descriptors(&self, descriptors: &Vec<vk::DescriptorSet>) {
        let njm_dev = self.dev.lock().unwrap();
        let r = unsafe { njm_dev.device().free_descriptor_sets(self.desc_pool, descriptors).unwrap() };
    }

    pub fn reset_pool(&self) {
        let njm_dev = self.dev.lock().unwrap();
        let r = unsafe { njm_dev.device().reset_descriptor_pool(self.desc_pool, vk::DescriptorPoolResetFlags::default()).unwrap() };
    }

}

impl Drop for NjmDescriptorPool {
    fn drop(&mut self) {
        let njm_dev = self.dev.lock().unwrap();
        unsafe { njm_dev.device().destroy_descriptor_pool(self.desc_pool, None) };
    }
}

// NjmDescriptorWriter
#[derive(Getters, Setters, WithSetters, MutGetters, CopyGetters, CloneGetters)]
pub struct NjmDescriptorWriter<'a> {
    #[getset(get = "pub")]
    set_layout: &'a NjmDescriptorSetLayout<'a>,

    #[getset(get = "pub")]
    pool: &'a NjmDescriptorPool,

    #[getset(get = "pub")]
    writes: Vec<vk::WriteDescriptorSet<'a>>,
}

impl<'a> NjmDescriptorWriter<'a> {
    pub fn new(set_layout: &'a NjmDescriptorSetLayout<'a>, pool: &'a NjmDescriptorPool) -> Self {
        return Self {
            set_layout,
            pool,
            writes: vec![],
        };
    }

    pub fn write_buffer(&mut self, binding: u32, buffer_info: &'a vk::DescriptorBufferInfo) -> &mut Self {
        let set_layout = self.set_layout;
        assert!(set_layout.bindings.contains_key(&binding), 
                "Layout does not contain specified binding");
        let binding_description = set_layout.bindings[&binding];
        assert!(binding_description.descriptor_count == 1, 
                "Binding single descriptor info, but binding expects multiple");
        
        let buffer_infos = std::slice::from_ref(buffer_info);
        let write = vk::WriteDescriptorSet::default()
                .descriptor_type(binding_description.descriptor_type)
                .descriptor_count(1)
                .dst_binding(binding)
                .buffer_info(buffer_infos);
        self.writes.push(write);
        return self;
    }

    pub fn write_image(&mut self, binding: u32, image_info: &'a vk::DescriptorImageInfo) -> &mut Self {
        let set_layout = self.set_layout;
        assert!(set_layout.bindings.contains_key(&binding), 
                "Layout does not contain specified binding");
        let binding_description = set_layout.bindings[&binding];
        assert!(binding_description.descriptor_count == 1, 
                "Binding single descriptor info, but binding expects multiple");
        
        let image_infos = std::slice::from_ref(image_info);
        let write = vk::WriteDescriptorSet::default()
                .descriptor_type(binding_description.descriptor_type)
                .descriptor_count(1)
                .dst_binding(binding)
                .image_info(image_infos);
        self.writes.push(write);
        return self;
    }

    pub fn build(&self) -> vk::DescriptorSet {
        let njm_desc_pool = self.pool;
        let njm_set_layout = self.set_layout;
        let desc_set = njm_desc_pool.allocate_descriptor(njm_set_layout.descriptor_set_layout);
        self.over_write(&desc_set);
        return desc_set;
    }

    pub fn over_write(&self, desc_set: &vk::DescriptorSet) {
        for write in &self.writes {
            write.dst_set(*desc_set);
        }
        let njm_pool = self.pool;
        let njm_dev = njm_pool.dev.lock().unwrap();
        unsafe { njm_dev.device().update_descriptor_sets(&self.writes, &[]) };
    }

}